{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "33e1519e-3d6f-4d80-9f54-3327ae98a5e6",
   "metadata": {},
   "source": [
    "# pykan Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a4eb37-b707-477a-bc7a-cd69bc730e0e",
   "metadata": {},
   "source": [
    "Some preliminaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e445797-b7e5-4245-aed6-7c8768bec575",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import autograd\n",
    "from tqdm import tqdm\n",
    "from kan import KAN as pykan\n",
    "\n",
    "import time\n",
    "from scipy.stats.qmc import Sobol\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.interpolate import RegularGridInterpolator\n",
    "import scipy\n",
    "device = torch.device(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5a0b1c6-720a-4b18-818d-4df1caa02ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Used for first-order derivatives\n",
    "def batch_jacobian(func, x, create_graph=False):\n",
    "    def _func_sum(x):\n",
    "        return func(x).sum(dim=0)\n",
    "    return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1, 0, 2)\n",
    "\n",
    "# Define Sobol Sampling\n",
    "def sobol_colloc(X0, X1, N, seed=42):\n",
    "    \"\"\"\n",
    "    N - number of colloc points\n",
    "    \n",
    "    \"\"\"\n",
    "    dims = X0.shape[0]\n",
    "    sobol_sampler = Sobol(dims, scramble=True, seed=seed)\n",
    "    points = sobol_sampler.random_base2(int(np.log2(N)))\n",
    "    points = X0 + points * (X1 - X0)\n",
    "    \n",
    "    return points"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf6fea14-baf3-4293-90e2-cda55ae862c8",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Diffusion Equation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "413e9b36-bc47-47db-a6ce-ea9f14ab4e9e",
   "metadata": {},
   "source": [
    "### Collocation Points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e47ae6b6-1a74-448e-a065-bbd19391e233",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 2**12\n",
    "collocs = np.array(sobol_colloc(np.array([0,-1]), np.array([1,1]), N)) # (4096, 2)\n",
    "\n",
    "# Generate Collocation points for BCs\n",
    "N = 2**6\n",
    "\n",
    "BC1_colloc = np.array(sobol_colloc(np.array([0,-1]), np.array([0,1]), N)) # (64, 2)\n",
    "BC1_data =  np.sin(np.pi*BC1_colloc[:,1]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC2_colloc = np.array(sobol_colloc(np.array([0,-1]), np.array([1,-1]), N)) # (64, 2)\n",
    "BC2_data = np.zeros(BC2_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC3_colloc = np.array(sobol_colloc(np.array([0,1]), np.array([1,1]), N)) # (64, 2)\n",
    "BC3_data = np.zeros(BC3_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "# Get collocation points from jaxkan application\n",
    "tcollocs = torch.from_numpy(collocs).to(device)\n",
    "tcollocs.requires_grad = True\n",
    "\n",
    "# Get boundary condition points and data from jaxkan application\n",
    "tbc_collocs = [torch.from_numpy(BC1_colloc).to(device),\n",
    "              torch.from_numpy(BC2_colloc).to(device),\n",
    "              torch.from_numpy(BC3_colloc).to(device)]\n",
    "\n",
    "tbc_data = [torch.from_numpy(BC1_data).to(device),\n",
    "           torch.from_numpy(BC2_data).to(device),\n",
    "           torch.from_numpy(BC3_data).to(device)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e43010ef-e4d4-44d6-9cdc-bb729b3ccae6",
   "metadata": {},
   "source": [
    "### Loss Function & training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cfe2a54-0bfc-4790-a3ca-b1432ad8a54b",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_vals = [3] # Grid sizes to use\n",
    "init_lr = 0.001\n",
    "num_epochs = 50000\n",
    "\n",
    "t_losses = []\n",
    "\n",
    "tstart_time = time.time()\n",
    "for grid in grid_vals:\n",
    "    if grid == grid_vals[0]:\n",
    "        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)\n",
    "        model.update_grid_from_samples(tcollocs)\n",
    "        model = model.to(device)\n",
    "    else:\n",
    "        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)\n",
    "        model = model.to(device)\n",
    "\n",
    "    # Define the training loop\n",
    "    def train():\n",
    "    \n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)\n",
    "\n",
    "        # Parameters\n",
    "        D = torch.tensor(0.15, device=device, requires_grad=False)\n",
    "        w_pde = torch.tensor(1.0, device=device, requires_grad=False)\n",
    "        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False)]\n",
    "\n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # The kan output has shape (batch, kan_output)\n",
    "            u = model(tcollocs)\n",
    "\n",
    "            # Define u_t and u_x functions\n",
    "            u_t_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)\n",
    "            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Define u_xx function\n",
    "            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Get u_t and u_x with shapes (batch,1)\n",
    "            u_t = u_t_fun(tcollocs)\n",
    "            #u_x = u_x_fun(tcollocs)\n",
    "            # Get u_xx with shape (batch,1)\n",
    "            u_xx = u_xx_fun(tcollocs)\n",
    "\n",
    "            # PDE Loss\n",
    "            source=torch.exp(-tcollocs[:,[0]])*(-torch.sin(torch.pi*tcollocs[:,[1]])+torch.pi**2*torch.sin(torch.pi*tcollocs[:,[1]]))\n",
    "            pde_loss = w_pde*torch.mean((u_t - u_xx-source)**2)\n",
    "            \n",
    "            # Boundary Conditions Loss\n",
    "            bc_loss = 0.0\n",
    "            for idx, _ in enumerate(tbc_collocs):\n",
    "                # Get model's prediction\n",
    "                pred = model(tbc_collocs[idx])\n",
    "                # Retrieve true values from boundary conditions\n",
    "                true = tbc_data[idx]\n",
    "                # MSE Loss\n",
    "                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)\n",
    "\n",
    "            total_loss = pde_loss + bc_loss\n",
    "            total_loss.backward()\n",
    "            return total_loss\n",
    "            \n",
    "        for step in range(num_epochs):\n",
    "    \n",
    "            optimizer.step(closure)\n",
    "\n",
    "            # Get Loss\n",
    "            current_loss = closure().item()\n",
    "            \n",
    "            t_losses.append(current_loss)\n",
    "    \n",
    "    train()\n",
    "\n",
    "tend_time = time.time()\n",
    "telapsed = tend_time - tstart_time\n",
    "print(f\"Total Time: {telapsed} s\")\n",
    "print(f\"Average time per iteration: {telapsed/num_epochs:.2f} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3137787d-82fa-4767-a6e2-6ee5c0aba5a3",
   "metadata": {},
   "source": [
    "### Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8d7cf08-93d9-440e-9c0d-c3357d53cd31",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "plt.plot(np.array(t_losses), label='Train Loss', marker='o', color='blue', markersize=1)\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training Loss Over Epochs')\n",
    "plt.yscale('log')  # Set y-axis to logarithmic scale\n",
    "\n",
    "plt.legend()\n",
    "plt.grid(True, which='both', linestyle='--', linewidth=0.5) \n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "827575f2-8be4-47a4-864d-11333fdcdc6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_t, N_x = 100, 256\n",
    "\n",
    "tt = torch.linspace(0, 1.0, N_t, device='cpu', requires_grad=False)\n",
    "tx = torch.linspace(0.0, 1.0, N_x, device='cpu', requires_grad=False)\n",
    "tT, tX = torch.meshgrid(tt, tx, indexing='ij')\n",
    "tcoords = torch.stack([tT.flatten(), tX.flatten()], dim=1).double().to(device)\n",
    "\n",
    "tresplot = model(tcoords)[:,0].cpu().detach().reshape(N_t, N_x)\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, tresplot, shading='auto', cmap='Spectral_r')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('Solution of Diffusion Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a1679fb-2b81-47f0-bbd2-cdb76c246592",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'pykan models/eq2-pykan')\n",
    "np.savez('../Plots/data/eq2-pykan.npz', arr_0=tresplot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93463341-d321-4faa-be78-05a4bf5f505d",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Helmholtz Equation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c42edc0-8cee-4638-9e6a-269758353950",
   "metadata": {},
   "source": [
    "### Collocation Points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85f13bc5-fab2-4af7-82aa-a2c279bd01ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate Collocation points for PDE\n",
    "N = 2**12\n",
    "collocs = np.array(sobol_colloc(np.array([-1,-1]), np.array([1,1]), N)) # (4096, 2)\n",
    "\n",
    "# Generate Collocation points for BCs\n",
    "N = 2**6\n",
    "\n",
    "BC1_colloc = np.array(sobol_colloc(np.array([-1,-1]), np.array([-1,1]), N)) # (64, 2)\n",
    "BC1_data =  np.zeros(BC1_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC2_colloc = np.array(sobol_colloc(np.array([-1,-1]), np.array([1,-1]), N)) # (64, 2)\n",
    "BC2_data = np.zeros(BC2_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC3_colloc = np.array(sobol_colloc(np.array([-1,1]), np.array([1,1]), N)) # (64, 2)\n",
    "BC3_data = np.zeros(BC3_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC4_colloc = np.array(sobol_colloc(np.array([1,-1]), np.array([1,1]), N)) # (64, 2)\n",
    "BC4_data = np.zeros(BC4_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "# Get collocation points from jaxkan application\n",
    "tcollocs = torch.from_numpy(collocs).to(device)\n",
    "tcollocs.requires_grad = True\n",
    "\n",
    "# Get boundary condition points and data from jaxkan application\n",
    "tbc_collocs = [torch.from_numpy(BC1_colloc).to(device),\n",
    "              torch.from_numpy(BC2_colloc).to(device),\n",
    "              torch.from_numpy(BC3_colloc).to(device),\n",
    "              torch.from_numpy(BC4_colloc).to(device)]\n",
    "\n",
    "tbc_data = [torch.from_numpy(BC1_data).to(device),\n",
    "           torch.from_numpy(BC2_data).to(device),\n",
    "           torch.from_numpy(BC3_data).to(device),\n",
    "           torch.from_numpy(BC4_data).to(device)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "855c4c59-6e8b-4b47-9fae-b4ef0704e545",
   "metadata": {},
   "source": [
    "### Loss Function & Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0942830-1dd8-44ff-9820-bcef44454211",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_vals = [3] # Grid sizes to use\n",
    "init_lr = 0.001\n",
    "num_epochs = 50000\n",
    "\n",
    "t_losses = []\n",
    "\n",
    "tstart_time = time.time()\n",
    "for grid in grid_vals:\n",
    "    if grid == grid_vals[0]:\n",
    "        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)\n",
    "        model.update_grid_from_samples(tcollocs)\n",
    "        model = model.to(device)\n",
    "    else:\n",
    "        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)\n",
    "        model = model.to(device)\n",
    "\n",
    "    # Define the training loop\n",
    "    def train():\n",
    "    \n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)\n",
    "\n",
    "        # Parameters\n",
    "        a1 = torch.tensor(1.0, device=device, requires_grad=False)\n",
    "        a2 = torch.tensor(4.0, device=device, requires_grad=False)\n",
    "        k = torch.tensor(1.0, device=device, requires_grad=False)\n",
    "        w_pde = torch.tensor(1.0, device=device, requires_grad=False)\n",
    "        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False)]\n",
    "\n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # The kan output has shape (batch, kan_output)\n",
    "            u = model(tcollocs)\n",
    "\n",
    "            # Define u_t and u_x functions\n",
    "            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)\n",
    "            u_y_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Define u_xx function\n",
    "            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,0].unsqueeze(1)\n",
    "            u_yy_fun = lambda x: batch_jacobian(u_y_fun, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Get u_t and u_x with shapes (batch,1)\n",
    "            u_xx = u_xx_fun(tcollocs)\n",
    "            #u_x = u_x_fun(tcollocs)\n",
    "            # Get u_xx with shape (batch,1)\n",
    "            u_yy = u_yy_fun(tcollocs)\n",
    "\n",
    "            # PDE Loss\n",
    "            sina1sina2=torch.sin(a1*torch.pi*tcollocs[:,[0]])*torch.sin(a2*torch.pi*tcollocs[:,[1]])\n",
    "            q= -(a1*torch.pi)**2*sina1sina2 -(a2*torch.pi)**2*sina1sina2 +k*sina1sina2\n",
    "            \n",
    "            pde_loss = w_pde*torch.mean((u_xx + u_yy +k**2 *u - q)**2)\n",
    "            \n",
    "            # Boundary Conditions Loss\n",
    "            bc_loss = 0.0\n",
    "            for idx, _ in enumerate(tbc_collocs):\n",
    "                # Get model's prediction\n",
    "                pred = model(tbc_collocs[idx])\n",
    "                # Retrieve true values from boundary conditions\n",
    "                true = tbc_data[idx]\n",
    "                # MSE Loss\n",
    "                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)\n",
    "\n",
    "            total_loss = pde_loss + bc_loss\n",
    "            total_loss.backward()\n",
    "            return total_loss\n",
    "            \n",
    "        for step in range(num_epochs):\n",
    "    \n",
    "            optimizer.step(closure)\n",
    "\n",
    "            # Get Loss\n",
    "            current_loss = closure().item()\n",
    "            \n",
    "            t_losses.append(current_loss)\n",
    "    \n",
    "    train()\n",
    "\n",
    "tend_time = time.time()\n",
    "telapsed = tend_time - tstart_time\n",
    "print(f\"Total Time: {telapsed} s\")\n",
    "print(f\"Average time per iteration: {telapsed/num_epochs:.2f} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3455aec1-52e8-4ba9-b7f3-d5378b170845",
   "metadata": {},
   "source": [
    "### Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd9b6c6b-3399-441e-9e87-d30b778bc01b",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_t, N_x = 100, 256\n",
    "\n",
    "tt = torch.linspace(0, 1.0, N_t, device='cpu', requires_grad=False)\n",
    "tx = torch.linspace(0.0, 1.0, N_x, device='cpu', requires_grad=False)\n",
    "tT, tX = torch.meshgrid(tt, tx, indexing='ij')\n",
    "tcoords = torch.stack([tT.flatten(), tX.flatten()], dim=1).double().to(device)\n",
    "\n",
    "tresplot = model(tcoords)[:,0].cpu().detach().reshape(N_t, N_x)\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, tresplot, shading='auto', cmap='Spectral_r')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('Solution of helmholtz Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2518239d-a89d-4fe2-bc7e-7323a40afd4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "a1=1.0\n",
    "a2=4.0\n",
    "k=1.0\n",
    "\n",
    "def u_exact(t,x): \n",
    "    return np.sin(a1*np.pi*t)*np.sin(a2*np.pi*x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9653c0af-6422-47db-bd92-532fdfd80955",
   "metadata": {},
   "outputs": [],
   "source": [
    "residuals = tresplot - u_exact(tT, tX)\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, residuals, shading='auto', cmap='seismic')#,vmin=-100,vmax=100)\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('residuals of Diffusion Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c2df11-1f7e-46af-b940-b0366590d980",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'pykan models/eq2-pykan')\n",
    "np.savez('../Plots/data/eq2-pykan.npz', arr_0=tresplot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c092aa2-a8b1-468c-b262-3441b597eaea",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Burgers Equation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e98d07da-4229-428d-8f09-3e4d141f7499",
   "metadata": {},
   "source": [
    "## Collocation Points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4532746b-d446-422b-abb7-7a9e0e0b7c1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate Collocation points for PDE\n",
    "N = 2**12\n",
    "collocs = np.array(sobol_colloc(np.array([-1,-1]), np.array([1,1]), N)) # (4096, 2)\n",
    "\n",
    "# Generate Collocation points for BCs\n",
    "N = 2**6\n",
    "\n",
    "BC1_colloc = np.array(sobol_colloc(np.array([0,-1]), np.array([0,1]), N)) # (64, 2)\n",
    "BC1_data = - np.sin(np.pi*BC1_colloc[:,1]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC2_colloc = np.array(sobol_colloc(np.array([0,-1]), np.array([1,-1]), N)) # (64, 2)\n",
    "BC2_data = np.zeros(BC2_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC3_colloc = np.array(sobol_colloc(np.array([0,1]), np.array([1,1]), N)) # (64, 2)\n",
    "BC3_data = np.zeros(BC3_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "# Get collocation points from jaxkan application\n",
    "tcollocs = torch.from_numpy(collocs).to(device)\n",
    "tcollocs.requires_grad = True\n",
    "\n",
    "# Get boundary condition points and data from jaxkan application\n",
    "tbc_collocs = [torch.from_numpy(BC1_colloc).to(device),\n",
    "              torch.from_numpy(BC2_colloc).to(device),\n",
    "              torch.from_numpy(BC3_colloc).to(device)]\n",
    "\n",
    "tbc_data = [torch.from_numpy(BC1_data).to(device),\n",
    "           torch.from_numpy(BC2_data).to(device),\n",
    "           torch.from_numpy(BC3_data).to(device)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53101445-87d3-4bfa-ad54-bac7813a98ae",
   "metadata": {},
   "source": [
    "## Loss function & training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a88c8318-9b27-4bbc-a8e5-1b2b89f00643",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_vals = [3] # Grid sizes to use\n",
    "init_lr = 0.001\n",
    "num_epochs = 50000\n",
    "\n",
    "t_losses = []\n",
    "\n",
    "tstart_time = time.time()\n",
    "for grid in grid_vals:\n",
    "    if grid == grid_vals[0]:\n",
    "        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)\n",
    "        model.update_grid_from_samples(tcollocs)\n",
    "        model = model.to(device)\n",
    "    else:\n",
    "        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)\n",
    "        model = model.to(device)\n",
    "\n",
    "    # Define the training loop\n",
    "    def train():\n",
    "    \n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)\n",
    "\n",
    "        # Parameters\n",
    "        v = torch.tensor(0.01/torch.pi, device=device, requires_grad=False)\n",
    "        w_pde = torch.tensor(1.0, device=device, requires_grad=False)\n",
    "        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False)]\n",
    "\n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # The kan output has shape (batch, kan_output)\n",
    "            u = model(tcollocs)\n",
    "\n",
    "            # Define u_t and u_x functions\n",
    "            u_t_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)\n",
    "            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Define u_xx function\n",
    "            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Get u_t and u_x with shapes (batch,1)\n",
    "            u_t = u_t_fun(tcollocs)\n",
    "            u_x = u_x_fun(tcollocs)\n",
    "            # Get u_xx with shape (batch,1)\n",
    "            u_xx = u_xx_fun(tcollocs)\n",
    "\n",
    "            # PDE Loss\n",
    "            pde_loss = w_pde*torch.mean((u_t + u*u_x - v*u_xx)**2)\n",
    "            \n",
    "            # Boundary Conditions Loss\n",
    "            bc_loss = 0.0\n",
    "            for idx, _ in enumerate(tbc_collocs):\n",
    "                # Get model's prediction\n",
    "                pred = model(tbc_collocs[idx])\n",
    "                # Retrieve true values from boundary conditions\n",
    "                true = tbc_data[idx]\n",
    "                # MSE Loss\n",
    "                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)\n",
    "\n",
    "            total_loss = pde_loss + bc_loss\n",
    "            total_loss.backward()\n",
    "            return total_loss\n",
    "            \n",
    "        for step in range(num_epochs):\n",
    "    \n",
    "            optimizer.step(closure)\n",
    "\n",
    "            # Get Loss\n",
    "            current_loss = closure().item()\n",
    "            \n",
    "            t_losses.append(current_loss)\n",
    "    \n",
    "    train()\n",
    "\n",
    "tend_time = time.time()\n",
    "telapsed = tend_time - tstart_time\n",
    "print(f\"Total Time: {telapsed} s\")\n",
    "print(f\"Average time per iteration: {telapsed/num_epochs:.2f} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa61b7ab-085a-4428-9231-c991558007ad",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63a6f777-66ca-457e-8b18-2b38a1bf6758",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "plt.plot(np.array(t_losses), label='Train Loss', marker='o', color='blue', markersize=1)\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training Loss Over Epochs')\n",
    "plt.yscale('log')  # Set y-axis to logarithmic scale\n",
    "\n",
    "plt.legend()\n",
    "plt.grid(True, which='both', linestyle='--', linewidth=0.5) \n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "643c675e-f690-4449-927e-8bc0bb16a61d",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_t, N_x = 100, 256\n",
    "\n",
    "tt = torch.linspace(0, 1.0, N_t, device='cpu', requires_grad=False)\n",
    "tx = torch.linspace(-1.0, 1.0, N_x, device='cpu', requires_grad=False)\n",
    "tT, tX = torch.meshgrid(tt, tx, indexing='ij')\n",
    "tcoords = torch.stack([tT.flatten(), tX.flatten()], dim=1).double().to(device)\n",
    "\n",
    "tresplot = model(tcoords)[:,0].cpu().detach().reshape(N_t, N_x)\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, tresplot, shading='auto', cmap='Spectral_r')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('Solution of Burgers Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e711523a-0de4-4401-bf75-839d9e5ff0ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "Bsol=np.load('../External Data/Burgers.npz')\n",
    "Bsol_interp=RegularGridInterpolator((Bsol['t'][:,0], Bsol['x'][:,0]) ,Bsol['usol'].T,method='linear',bounds_error=False,fill_value=np.nan)\n",
    "\n",
    "def u_exact(t,x): return Bsol_interp((t, t))\n",
    "\n",
    "residuals = tresplot - u_exact(tT, tX)\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, residuals, shading='auto', cmap='seismic')#,vmin=-100,vmax=100)\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('Solution of Burgers Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c30937-0ef7-488b-9cc6-0d913697bcc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'pykan models/eq3-pykan')\n",
    "np.savez('../Plots/data/eq3-pykan.npz', arr_0=tresplot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8db00b6a-71be-4d5a-b9c1-a0d6f91d5c4b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "# Allen Cahn Equation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8a377ab-c850-405e-94c8-ae6239685b1c",
   "metadata": {},
   "source": [
    "## Collocation Points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "170b79ee-a68b-4ef9-b2d8-04a107d1922b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate Collocation points for PDE\n",
    "N = 2**12\n",
    "collocs = np.array(sobol_colloc(np.array([0,-1]), np.array([1,1]), N)) # (4096, 2)\n",
    "\n",
    "# Generate Collocation points for BCs\n",
    "N = 2**6\n",
    "\n",
    "BC1_colloc = np.array(sobol_colloc(np.array([0,-1]), np.array([0,1]), N)) # (64, 2)\n",
    "BC1_data =  ((BC1_colloc[:,1]**2)*np.cos(np.pi*BC1_colloc[:,1])).reshape(-1,1)\n",
    "\n",
    "BC2_colloc = np.array(sobol_colloc(np.array([0,-1]), np.array([1,-1]), N)) # (64, 2)\n",
    "BC2_data = -np.ones(BC2_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "BC3_colloc = np.array(sobol_colloc(np.array([0,1]), np.array([1,1]), N)) # (64, 2)\n",
    "BC3_data = -np.ones(BC3_colloc.shape[0]).reshape(-1,1) # (64, 1)\n",
    "\n",
    "# Get collocation points from jaxkan application\n",
    "tcollocs = torch.from_numpy(collocs).to(device)\n",
    "tcollocs.requires_grad = True\n",
    "\n",
    "# Get boundary condition points and data from jaxkan application\n",
    "tbc_collocs = [torch.from_numpy(BC1_colloc).to(device),\n",
    "              torch.from_numpy(BC2_colloc).to(device),\n",
    "              torch.from_numpy(BC3_colloc).to(device)]\n",
    "\n",
    "tbc_data = [torch.from_numpy(BC1_data).to(device),\n",
    "           torch.from_numpy(BC2_data).to(device),\n",
    "           torch.from_numpy(BC3_data).to(device)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c296ba6-ef88-44c4-8e78-f3d4e8ea3ac2",
   "metadata": {},
   "source": [
    "## Loss Function & training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cf2fb9b-74e6-4945-9182-08efd2335a84",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_vals = [3] # Grid sizes to use\n",
    "init_lr = 0.001\n",
    "num_epochs = 50000\n",
    "\n",
    "t_losses = []\n",
    "\n",
    "tstart_time = time.time()\n",
    "for grid in grid_vals:\n",
    "    if grid == grid_vals[0]:\n",
    "        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)\n",
    "        model.update_grid_from_samples(tcollocs)\n",
    "        model = model.to(device)\n",
    "    else:\n",
    "        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)\n",
    "        model = model.to(device)\n",
    "\n",
    "    # Define the training loop\n",
    "    def train():\n",
    "    \n",
    "        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)\n",
    "\n",
    "        # Parameters\n",
    "        D = torch.tensor(0.001, device=device, requires_grad=False)\n",
    "        c = torch.tensor(5.0, device=device, requires_grad=False)\n",
    "        w_pde = torch.tensor(1.0, device=device, requires_grad=False)\n",
    "        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False),\n",
    "                 torch.tensor(1.0, device=device, requires_grad=False)]\n",
    "\n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # The kan output has shape (batch, kan_output)\n",
    "            u = model(tcollocs)\n",
    "\n",
    "            # Define u_t and u_x functions\n",
    "            u_t_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)\n",
    "            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Define u_xx function\n",
    "            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,1].unsqueeze(1)\n",
    "\n",
    "            # Get u_t and u_x with shapes (batch,1)\n",
    "            u_t = u_t_fun(tcollocs)\n",
    "            #u_x = u_x_fun(tcollocs)\n",
    "            # Get u_xx with shape (batch,1)\n",
    "            u_xx = u_xx_fun(tcollocs)\n",
    "\n",
    "            # PDE Loss\n",
    "            pde_loss = w_pde*torch.mean((u_t - D*u_xx-c*(u-u**3))**2)\n",
    "            \n",
    "            # Boundary Conditions Loss\n",
    "            bc_loss = 0.0\n",
    "            for idx, _ in enumerate(tbc_collocs):\n",
    "                # Get model's prediction\n",
    "                pred = model(tbc_collocs[idx])\n",
    "                # Retrieve true values from boundary conditions\n",
    "                true = tbc_data[idx]\n",
    "                # MSE Loss\n",
    "                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)\n",
    "\n",
    "            total_loss = pde_loss + bc_loss\n",
    "            total_loss.backward()\n",
    "            return total_loss\n",
    "            \n",
    "        for step in range(num_epochs):\n",
    "    \n",
    "            optimizer.step(closure)\n",
    "\n",
    "            # Get Loss\n",
    "            current_loss = closure().item()\n",
    "            \n",
    "            t_losses.append(current_loss)\n",
    "    \n",
    "    train()\n",
    "\n",
    "tend_time = time.time()\n",
    "telapsed = tend_time - tstart_time\n",
    "print(f\"Total Time: {telapsed} s\")\n",
    "print(f\"Average time per iteration: {telapsed/num_epochs:.2f} s\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6228fa47-eb26-48f5-9981-3fe4aa836b39",
   "metadata": {},
   "source": [
    "## Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2e0b175-5fbc-4942-b6ec-7110b6f7434a",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 6))\n",
    "\n",
    "plt.plot(np.array(t_losses), label='Train Loss', marker='o', color='blue', markersize=1)\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.title('Training Loss Over Epochs')\n",
    "plt.yscale('log')  # Set y-axis to logarithmic scale\n",
    "\n",
    "plt.legend()\n",
    "plt.grid(True, which='both', linestyle='--', linewidth=0.5) \n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f4fecef-c951-499b-9b20-54090ebef05a",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_t, N_x = 100, 256\n",
    "\n",
    "tt = torch.linspace(0, 1.0, N_t, device='cpu', requires_grad=False)\n",
    "tx = torch.linspace(-1.0, 1.0, N_x, device='cpu', requires_grad=False)\n",
    "tT, tX = torch.meshgrid(tt, tx, indexing='ij')\n",
    "tcoords = torch.stack([tT.flatten(), tX.flatten()], dim=1).double().to(device)\n",
    "\n",
    "tresplot = model(tcoords)[:,0].cpu().detach().reshape(N_t, N_x)\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, tresplot, shading='auto', cmap='Spectral_r')\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('Solution of Allen Cahn Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81b6813d-7682-4cf8-8782-4f06e8f727e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw reference values\n",
    "ref = scipy.io.loadmat('../External Data/usol_D_0.001_k_5.mat')\n",
    "\n",
    "sol_interp=RegularGridInterpolator((ref['t'][0], ref['x'][0]) ,ref['u'],method='linear',bounds_error=False,fill_value=np.nan)\n",
    "\n",
    "def u_exact(t,x): return sol_interp((t, x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fe17924-732e-45c2-bd84-4695bb9d509b",
   "metadata": {},
   "outputs": [],
   "source": [
    "residuals = tresplot - u_exact(tT, tX)\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.pcolormesh(tT, tX, residuals, shading='auto', cmap='seismic')#,vmin=-100,vmax=100)\n",
    "plt.colorbar()\n",
    "\n",
    "plt.title('Solution of Allen Cahn Equation')\n",
    "plt.xlabel('t')\n",
    "\n",
    "plt.ylabel('x')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e716e772-f164-492b-aac2-392da16d8a5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'pykan models/eq4-pykan')\n",
    "np.savez('../Plots/data/eq4-pykan.npz', arr_0=tresplot)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
